import pdb
import os
import pandas as pd
from dataset_modules.dataset_generic import Generic_WSI_Classification_Dataset, Generic_MIL_Dataset, save_splits
import argparse
import numpy as np

parser = argparse.ArgumentParser(description='Creating splits for whole slide classification')
parser.add_argument('--label_frac', type=float, default= 1.0,
                    help='fraction of labels (default: 1)')
parser.add_argument('--seed', type=int, default=1,
                    help='random seed (default: 1)')
parser.add_argument('--k', type=int, default=10,
                    help='number of splits (default: 10)')
parser.add_argument('--task', type=str, choices=['MSI_TCGA','MSI_PAIP', 'MSI_HCH', 'MSI_synthetic_HCH','camelyon','camelyon16_0.5','MSI_HCH_0.6','TCGA_NSCLC'])
parser.add_argument('--val_frac', type=float, default= 0.1,
                    help='fraction of labels for validation (default: 0.1)')
parser.add_argument('--test_frac', type=float, default= 0.1,
                    help='fraction of labels for test (default: 0.1)')

args = parser.parse_args()

if args.task == 'task_1_tumor_vs_normal':
    args.n_classes=2
    dataset = Generic_WSI_Classification_Dataset(csv_path = 'dataset_csv/tumor_vs_normal_dummy_clean.csv',
                            shuffle = False, 
                            seed = args.seed, 
                            print_info = True,
                            label_dict = {'normal_tissue':0, 'tumor_tissue':1},
                            patient_strat=True,
                            ignore=[])

elif args.task == 'task_2_tumor_subtyping':
    args.n_classes=3
    dataset = Generic_WSI_Classification_Dataset(csv_path = 'dataset_csv/tumor_subtyping_dummy_clean.csv',
                            shuffle = False, 
                            seed = args.seed, 
                            print_info = True,
                            label_dict = {'subtype_1':0, 'subtype_2':1, 'subtype_3':2},
                            patient_strat= True,
                            patient_voting='maj',
                            ignore=[])
elif args.task == 'MSI_PAIP':
    args.n_classes = 2
    dataset = Generic_WSI_Classification_Dataset(csv_path='dataset_csv/MSI_classification_PAIP.csv',
                                                 shuffle=False,
                                                 seed=args.seed,
                                                 print_info=False,
                                                 label_dict={'MSS': 0, 'MSI-H': 1},
                                                 patient_strat=True,
                                                 ignore=[])
elif args.task == 'MSI_TCGA':
    args.n_classes = 2
    dataset = Generic_WSI_Classification_Dataset(csv_path='dataset_csv/MSI_classification_TCGA_S1e9+.csv',
                                                 shuffle=False,
                                                 seed=args.seed,
                                                 print_info=False,
                                                 label_dict={'MSS': 0, 'MSI-H': 1},
                                                 patient_strat=True,
                                                 ignore=[])
elif args.task == 'MSI_HCH':
    args.n_classes = 3
    dataset = Generic_WSI_Classification_Dataset(csv_path='dataset_csv/MSI_classification_HCH20221202.csv',
                                                 shuffle=False,
                                                 seed=args.seed,
                                                 print_info=False,
                                                 label_dict={'MSS': 0, 'MSI-H': 1, 'Normal': 2},
                                                 patient_strat=True,
                                                 ignore=[])
elif args.task == 'TCGA_NSCLC':
    args.n_classes = 3
    dataset = Generic_WSI_Classification_Dataset(csv_path='/mnt/sda2/WSI/TCGA-NSCLC/TCGA-NSCLC.csv',
                                                 shuffle=False,
                                                 seed=args.seed,
                                                 print_info=False,
                                                 label_dict={'LUAD': 0, 'LUSC': 1},
                                                 patient_strat=True,
                                                 ignore=[])
elif args.task == 'MSI_HCH_0.6':
    args.n_classes = 3
    dataset = Generic_WSI_Classification_Dataset(csv_path='dataset_csv/MSI_classification_HCH20221202.csv',
                                                 shuffle=False,
                                                 seed=args.seed,
                                                 print_info=False,
                                                 label_dict={'MSS': 0, 'MSI-H': 1, 'Normal': 2},
                                                 patient_strat=True,
                                                 ignore=[])

elif args.task == 'camelyon':
    args.n_classes = 2
    dataset = Generic_WSI_Classification_Dataset(csv_path='dataset_csv/filtered_camelyon_data.csv',
                                                 shuffle=False,
                                                 seed=args.seed,
                                                 print_info=True,
                                                 label_dict={'normal_tissue': 0, 'tumor_tissue': 1},
                                                 patient_strat=True,
                                                 ignore=[])
elif args.task == 'camelyon16_0.5':
    args.n_classes = 2
    dataset = Generic_WSI_Classification_Dataset(csv_path='dataset_csv/camelyon16_add.csv',
                                                 shuffle=False,
                                                 seed=args.seed,
                                                 print_info=True,
                                                 label_dict={'normal': 0, 'tumor': 1},
                                                 patient_strat=True,
                                                 ignore=[])

else:
    raise NotImplementedError
# patient_cls_ids记录了各个类包含哪些病人
num_slides_cls = np.array([len(cls_ids) for cls_ids in dataset.patient_cls_ids])
val_num = np.round(num_slides_cls * args.val_frac).astype(int)
test_num = np.round(num_slides_cls * args.test_frac).astype(int)

if __name__ == '__main__':
    if args.label_frac > 0:
        label_fracs = [args.label_frac]
    else:
        label_fracs = [0.1, 0.25, 0.5, 0.75, 1.0]
    #  在整个数据集中，随机选择占比label_frac的数据，生成子数据集，然后再在里面分成 train,test,val
    for lf in label_fracs:
        split_dir = 'splits/'+ str(args.task) + '_{}'.format(int(lf * 100))
        os.makedirs(split_dir, exist_ok=True)
        dataset.create_splits(k = args.k, val_num = val_num, test_num = test_num, label_frac=lf) # 返回的是一个分割的迭代器
        for i in range(args.k):
            dataset.set_splits() # 病人id集 转到 切片id集
            descriptor_df = dataset.test_split_gen(return_descriptor=True)
            splits = dataset.return_splits(from_id=True)
            save_splits(splits, ['train', 'val', 'test'], os.path.join(split_dir, 'splits_{}.csv'.format(i)))
            save_splits(splits, ['train', 'val', 'test'], os.path.join(split_dir, 'splits_{}_bool.csv'.format(i)), boolean_style=True)
            # save_splits(splits, ['train', 'train_label', 'val', 'val_label', 'test', 'test_label'],
            #             os.path.join(split_dir, 'splits_{}.csv'.format(i)))
            # save_splits(splits, ['train', 'train_label', 'val', 'val_label', 'test', 'test_label'],
            #             os.path.join(split_dir, 'splits_{}_bool.csv'.format(i)), boolean_style=True)
            descriptor_df.to_csv(os.path.join(split_dir, 'splits_{}_descriptor.csv'.format(i)))



